import subprocess
import argparse
import asyncio
import aiohttp
import sys
import os
import requests
import fitz
import configs

nougat_conf = configs.Nougat()

def _get_mmd_path(pdf_path, mmd_dir):
    filename, _ = os.path.splitext(os.path.basename(pdf_path))
    mmd_path = f'{mmd_dir}/{filename}.mmd'
    return mmd_path

def _parse_online(pdf_path, mmd_dir, rm_pdf, stop=None):
    mmd_path = None
    with open(pdf_path, 'rb') as fin:
        headers = {'accept': 'application/json'}
        files = {'file': (pdf_path, fin, 'application/pdf')}
        params = f'?start=1&stop={stop}' if stop else ''
        response = requests.post(f'{nougat_conf.base_url}{params}', headers=headers, files=files)
        if response.status_code == 200:
            mmd_path = _get_mmd_path(pdf_path, mmd_dir)
            with open(mmd_path, 'w') as fout:
                print(response.json(), file=fout)
        else:
            print(f'Requesting nougat api failed, error code: {response.status_code}，Info: {response.text}',
                  file=sys.stderr)

    if rm_pdf:
        try:
            command_rm = ['/bin/rm', pdf_path]
            result_rm = subprocess.run(command_rm, capture_output=True, text=True)
        except FileNotFoundError as e:
            print(f'The command "/bin/rm" was not found: {e}')

    return mmd_path

def run_parse(pdf_path, mmd_dir, rm_pdf=False, stop=None):
    mmd_path = _parse_online(pdf_path, mmd_dir, rm_pdf=rm_pdf, stop=stop)

    page_path = None
    if mmd_path:
        try:
            doc = fitz.open(pdf_path)
            page = doc.load_page(0)
            pix = page.get_pixmap()
            page_path = os.path.splitext(pdf_path)[0] + '_page.png'
            pix.save(page_path)
            doc.close()
        except Exception as e:
            raise f'Error in processing PDF file: {str(e)}'

    return mmd_path, page_path

if __name__=='__main__':
    mmd_path, page_path = run_parse('_data/test.pdf', '_data/')
    print(mmd_path, page_path)
